package com.jojo.rateLimiter;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * 漏桶算法
 *
 * @author hpl
 * @since 2024/12/16
 */
public class LeakyBucket {
    // 桶的最大容量
    private final int capacity;
    // 每秒流出的速率
    private final int leakRate;
    // 当前桶中的水（请求数）
    private AtomicInteger currentWater = new AtomicInteger(0);
    // 上次漏水时间
    private long lastLeakTime;

    public LeakyBucket(int capacity, int leakRate) {
        this.capacity = capacity;
        this.leakRate = leakRate;
        this.lastLeakTime = System.currentTimeMillis();
    }

    /**
     * 尝试放入一个请求
     *
     * @return 如果请求成功放入桶中返回 true；否则返回 false
     */
    public synchronized boolean tryAdd() {
        // 计算漏水后桶中的剩余水量
        leak();

        // 如果桶未满，则放入请求
        if (currentWater.get() < capacity) {
            currentWater.incrementAndGet();
            return true;
        } else {
            // 桶已满，请求被丢弃
            return false;
        }
    }

    /**
     * 根据时间计算需要漏掉的水
     */
    private void leak() {
        long now = System.currentTimeMillis();
        long elapsedTime = now - lastLeakTime;

        // 计算应漏掉的水量
        int leaked = (int) (elapsedTime / 1000.0 * leakRate);
        if (leaked > 0) {
            // 更新桶中的水量
            currentWater.updateAndGet(water -> Math.max(0, water - leaked));
            // 更新上次漏水时间
            lastLeakTime = now;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        // 桶的容量为 10，每秒漏出 2 个请求
        LeakyBucket bucket = new LeakyBucket(10, 2);

        // 模拟请求
        for (int i = 0; i < 20; i++) {
            if (bucket.tryAdd()) {
                System.out.println("第 " + (i + 1) + " 个请求成功放入桶中");
            } else {
                System.out.println("第 " + (i + 1) + " 个请求被丢弃");
            }
            Thread.sleep(300); // 每隔 300 毫秒发起一个请求
        }
    }
}
